#!/usr/bin/python

from bcc import BPF
import sys
import socket
import struct
import time

def usage():
    print("Usage: {0} <ifdev> ".format(sys.argv[0]))
    print("e.g.: {0} enp2s0 (default enp2s0) \n".format(sys.argv[0]))

if len(sys.argv) != 2:
    usage()
    ifdev = "enp2s0"
else:
    ifdev = sys.argv[1]

# load BPF program
b = BPF(text = """
#include <linux/bpf.h>
#include <linux/in.h>
#include <linux/if_ether.h>
#include <linux/ip.h>
#include <linux/types.h>
#include <stddef.h>
#include <memory.h>
#include <sys/types.h> 
        

BPF_LPM_TRIE(blacklist, u64, u32, 16);

int firewall(struct xdp_md *ctx) {
  void *data_end = (void *)(long)ctx->data_end;
  void *data = (void *)(long)ctx->data;

  // Only IPv4 supported for this example
  struct ethhdr *ether = data;
  if (data + sizeof(*ether) > data_end) {
    // Malformed Ethernet header
    return XDP_ABORTED;
  }

  if (ether->h_proto != 0x08U) {  // htons(ETH_P_IP) -> 0x08U
    // Non IPv4 traffic
    return XDP_PASS;
  }

  data += sizeof(*ether);
  struct iphdr *ip = data;
  if (data + sizeof(*ip) > data_end) {
    // Malformed IPv4 header
    return XDP_ABORTED;
  }

  struct {
    __u32 prefixlen;
    __u32 saddr;
  } key;

  key.prefixlen = 32;
  key.saddr = ip->saddr;

  // Lookup SRC IP in blacklisted IPs
  __u64 *rule_idx = blacklist.lookup(&key);
  if (rule_idx) {
    // Matched, increase match counter for matched "rule"
    __u32 index = *(__u32*)rule_idx;  // make verifier happy
    return XDP_DROP;
  }

  return XDP_PASS;
}
""", cflags=["-w"])
        

fn = b.load_func("firewall", BPF.XDP)
b.attach_xdp(ifdev, fn)

# 获取黑名单 map
blacklist = b["blacklist"]

# 添加 IP 地址到黑名单
ip = socket.inet_aton("192.168.1.1")  # 将 IP 地址转换为网络字节序
key = struct.pack("II", 32, struct.unpack("!I", ip)[0])
blacklist[key] = struct.pack("B", 1)

print("Blacklist loaded. Monitoring packets...")
while 1:
    try:
        time.sleep(1)
    except KeyboardInterrupt:
        print("Removing filter from device")
        break

b.remove_xdp(ifdev)
